// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/process/profiler.h>
#include <melon/futures/future.h>
#include <melon/futures/promise.h>
#include <melon/init/init.h>
#include <gtest/gtest.h>
#include <signal.h>
#include <thread>
#include <pollux/common/process/trace_context.h>

using namespace kumo::pollux::process;
using namespace kumo::pollux;

namespace {
int32_t fi(int32_t x) {
  return x < 2 ? x : fi(x - 1) + fi(x - 2);
}
void compute(int32_t seconds) {
  auto start = getCurrentTimeMs();
  constexpr int32_t kNumThreads = 10;
  for (;;) {
    std::vector<std::thread> threads;
    threads.reserve(kNumThreads);
    std::atomic<int32_t> sum = 0;
    for (int32_t i = 0; i < kNumThreads; ++i) {
      threads.push_back(std::thread([&]() {
        sum += fi(40);
        std::this_thread::sleep_for(std::chrono::milliseconds(3)); // NOLINT
      }));
    }
    for (auto& thread : threads) {
      thread.join();
    }
    KLOG(INFO) << "Sum " << sum;
    if (getCurrentTimeMs() - start > seconds * 1000) {
      break;
    }
  }
}

} // namespace

TEST(ProfilerTest, basic) {
#if !defined(linux)
  return;
#endif
  filesystems::registerLocalFileSystem();
  // We have seconds of busy and idle activity. We set the profiler to
  // check every second and to trigger after 1s at 200%. A burst of
  // under 2s is not recorded and a new file is started after every 4s
  // of cpu busy.

  turbo::set_flag(&FLAGS_profiler_check_interval_seconds, 1);
  turbo::set_flag(&FLAGS_profiler_min_cpu_pct , 200);
  turbo::set_flag(&FLAGS_profiler_max_sample_seconds , 4);
  turbo::set_flag(&FLAGS_profiler_max_sample_seconds , 2);

  Profiler::start("/tmp/profilertest");
  compute(5);
  std::this_thread::sleep_for(std::chrono::seconds(2)); // NOLINT
  compute(1);
  std::this_thread::sleep_for(std::chrono::seconds(2)); // NOLINT

  compute(3);
  Profiler::stop();

  // We set the profiler to start regardless of load and wait 30s before
  // producing the next result.
  turbo::set_flag(&FLAGS_profiler_check_interval_seconds , 30);
  turbo::set_flag(&FLAGS_profiler_min_cpu_pct , 0);
  turbo::set_flag(&FLAGS_profiler_min_sample_seconds , 0);
  Profiler::start("/tmp/profilertest");
  compute(2);
  // The test exits during the measurement interval. We expect no
  // crash on exit if the threads are properly joined.
}

int main(int argc, char** argv) {
  // Fork a child process to run all the tests.
  int32_t pid = fork();
  if (pid < 0) {
    KLOG(ERROR) << "Failed to fork child";
    exit(1);
  }
  if (pid > 0) {
    // The parent waits for the child to return. If the child returns
    // in time, the child's return code is returned. If the child does
    // not return in time, we return 0 and the test fails silently.
    std::atomic<bool> timedOut = false;
    std::atomic<bool> completed = false;
    auto sleepPromise = melon::Promise<bool>();
    melon::SemiFuture<bool> sleepFuture(false);
    sleepFuture = sleepPromise.getSemiFuture();
    std::thread timer([&]() {
      try {
        auto& executor = melon::QueuedImmediateExecutor::instance();
        // Wait for up to 100 seconds. The test is normally ~20s unless it
        // hangs.
        std::move(sleepFuture).via(&executor).wait((std::chrono::seconds(100)));
      } catch (std::exception&) {
      }
      if (completed) {
        return;
      }
      timedOut = true;
      KLOG(INFO) << "Killing the test process for timeout";
      kill(pid, SIGKILL);
    });

    int wstatus;
    waitpid(pid, &wstatus, WUNTRACED | WCONTINUED);
    KLOG(INFO) << "Test completed";
    completed = true;
    sleepPromise.setValue(true);
    timer.join();

    if (timedOut) {
      return 0;
    }
    return WEXITSTATUS(wstatus);
  }

  testing::InitGoogleTest(&argc, argv);
  // Signal handler required for ThreadDebugInfoTest
  melon::Init init(&argc, &argv, false);
  return RUN_ALL_TESTS();
  return 0;
}
